package com.ai.service.impl;


import com.ai.config.PlanetAIProperties;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.transformer.splitter.TextSplitter;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.InputStreamResource;
import org.springframework.core.io.Resource;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;

import java.io.IOException;
import java.util.List;


/**
 * 文档分割服务
 * 直接从 TikaDocumentReader 获取的 Document 对象可能包含大量文本，通常需要进一步分割处理，
 * 特别是为 RAG (检索增强生成) 应用准备数据时
 */
@Service
public class DocumentSplitterService {

    private final DocumentParserService tikaDocumentService;
    private final TokenTextSplitter tokenTextSplitter;

    @Autowired
    public DocumentSplitterService(DocumentParserService tikaDocumentService, PlanetAIProperties planetAIProperties) {
        this.tikaDocumentService = tikaDocumentService;
        // 配置文本分割器
        this.tokenTextSplitter = new TokenTextSplitter(
                planetAIProperties.getRag().getChunkSize(),
                planetAIProperties.getRag().getMinChunkSizeChars(),
                planetAIProperties.getRag().getMinChunkLengthToEmbed(),
                planetAIProperties.getRag().getMaxNumChunks(),
                planetAIProperties.getRag().getKeepSeparator()
        );
    }

    /**
     * 文档分割、自动分割
     * @param file MultipartFile
     * @return List
     * @throws IOException E
     */
    public List<String> splitDocumentContent(MultipartFile file) throws IOException {
        // 1. 使用 TikaDocumentReader 读取文档
        Resource resource = new InputStreamResource(file.getInputStream());
        TikaDocumentReader tikaDocumentReader = new TikaDocumentReader(resource);
        List<Document> originalDocuments = tikaDocumentReader.get();

        // 2. 初始化文本分割器（例如按 Token 数分割）
        TextSplitter textSplitter = new TokenTextSplitter(); // 可配置参数如块大小、重叠大小

        // 3. 应用分割器
        List<Document> splitDocuments = textSplitter.apply(originalDocuments);

        // 4. 提取分割后的文本内容并返回
        return splitDocuments.stream()
                .map(Document::getText)
                .toList();
    }


    /**
     * 分割上传的文件内容 自定义分割
     * @param file 上传的文件
     * @return 分割后的文本内容列表
     */
    public List<String> customerSplitDocumentContent(MultipartFile file) {
        try {
            // 1. 使用 TikaDocumentReader 读取文档
            List<Document> originalDocuments = tikaDocumentService.parseMultipartFileToDocuments(file);

            // 2. 应用分割器
            List<Document> splitDocuments = tokenTextSplitter.apply(originalDocuments);

            // 3. 提取分割后的文本内容并返回
            return splitDocuments.stream()
                    .map(Document::getText)
                    .toList();
        } catch (Exception e) {
            throw new RuntimeException("文件分割失败: " + e.getMessage(), e);
        }
    }

    /**
     * 分割文档并返回完整的 Document 列表（包含元数据） chunks
     * @param file 上传的文件
     * @return 分割后的 Document 列表
     */
    public List<Document> splitDocumentToDocuments(MultipartFile file) {
        try {
            List<Document> originalDocuments = tikaDocumentService.parseMultipartFileToDocuments(file);
            return tokenTextSplitter.apply(originalDocuments);
        } catch (Exception e) {
            throw new RuntimeException("文件分割失败: " + e.getMessage(), e);
        }
    }
}
